package _02_greedy_algorithm.huffmantree;

import java.util.*;

/**
 * @author: mornd
 * @dateTime: 2023/7/20 - 11:00
 * 哈夫曼树定义
 */
public class HuffmanTree {
    String str;
    Map<Character, Node> map = new HashMap<>();

    int bits;

    Node root;

    public HuffmanTree(String str) {
        this.str = str;
        char[] chars = str.toCharArray();
        for (char c : chars) {
//            if (!map.containsKey(c)) {
//                map.put(c, new Node(c));
//            }
//            Node node = map.get(c);
//            node.freq++;

            Node node = map.computeIfAbsent(c, Node::new);
            node.freq++;
        }

        PriorityQueue<Node> queue = new PriorityQueue<>(
                Comparator.comparingInt(Node::getFreq)
        );
        queue.addAll(map.values());

        while (queue.size() >= 2) {
            Node x = queue.poll();
            Node y = queue.poll();

            int freq = x.freq + y.freq;
            Node parent = new Node(freq, x, y);

            queue.offer(parent);
        }

        root = queue.poll();
        bits = dfs(root, new StringBuilder());
//        System.out.println(root);
    }

    // 深度遍历
    private int dfs(Node n, StringBuilder code) {
        int sum = 0;
        if (n.isLeaf()) {
            n.code = code.toString();
            sum = n.freq * n.code.length();
//            System.out.println(n + ":" + code);
        } else {
            sum += dfs(n.left, code.append("0"));
            sum += dfs(n.right, code.append("1"));
        }
        if (code.length() > 0) {
            code.deleteCharAt(code.length() - 1);
        }
        return sum;
    }

    // 编码
    public String encode() {
        char[] chars = str.toCharArray();
        StringBuilder sb = new StringBuilder();
        for (char c : chars) {
            String code = map.get(c).code;
            sb.append(code);
        }
        return sb.toString();
    }

    // 解码
    public String decode(String str) {
        char[] chars = str.toCharArray();
        StringBuilder sb = new StringBuilder();
        /*

                      7
                     /\
                    3  c_4
                   /\
                 a_1 b_2

                    0001011111
         */
        int i = 0;
        Node curr = root;
        while (i < chars.length) {
            if (!curr.isLeaf()) {
                // 默认0往左边找，1往右边找
                if (str.charAt(i) == '0') {
                    curr = curr.left;
                } else if ((str.charAt(i) == '1')) {
                    curr = curr.right;
                }
                i++;
            }
            if (curr.isLeaf()) {
                // 到叶子节点了，获取解码的值
                sb.append(curr.ch);
                // 当前节点重置为root
                curr = root;
            }
        }
        return sb.toString();
    }

    public void printCodeValue() {
        for (Node value : map.values()) {
            System.out.println(value.ch + "->" + value.code);
        }
    }

    // 满二叉树结构
    static class Node {
        // 保存的字符
        Character ch;
        // 出现频次
        int freq;
        Node left;
        Node right;
        // 编码
        String code;

        public Node(Character ch) {
            this.ch = ch;
        }

        public Node(int freq, Node left, Node right) {
            this.freq = freq;
            this.left = left;
            this.right = right;
        }

        public int getFreq() {
            return freq;
        }

        // 是否叶子节点，前提该树是满二叉树
        boolean isLeaf() {
            return left == null;
        }

        @Override
        public String toString() {
            return new StringJoiner(", ", Node.class.getSimpleName() + "[", "]")
                    .add("ch=" + ch)
                    .add("freq=" + freq)
                    .toString();
        }
    }
}
